#include <ros/ros.h>
#include <sensor_msgs/LaserScan.h>
#include <nav_msgs/Odometry.h>
#include <map.hpp>
#include <transform.h>
#include <Eigen/Dense>
#include <visualization_msgs/Marker.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>

class MapMatch{
public:
    MapMatch(ros::NodeHandle& n)
    {
        nh = n;
        std::string scan_topic,scan_pub_topic,odom_topic,map_topic,localization_result_topic;
        nh.getParam("scan_topic",scan_topic);
        nh.getParam("odom_topic",odom_topic);
        nh.getParam("localization_result_topic",localization_result_topic);
        nh.getParam("map_topic",map_topic);
        nh.getParam("iter_times",iter_times_);
        nh.getParam("sample_interval",sample_interval_);
        nh.getParam("scan_pub_topic",scan_pub_topic);
        scan_sub_ = nh.subscribe<sensor_msgs::LaserScan>(scan_topic,5,&MapMatch::scan_callback,this);
        odom_sub_ = nh.subscribe<nav_msgs::Odometry>(odom_topic,5,&MapMatch::odom_callback,this);
        relocalization_sub_ = nh.subscribe<geometry_msgs::PoseWithCovarianceStamped>("/initialpose",5,&MapMatch::reloc_callback,this);
        test_scan_pub_ = nh.advertise<sensor_msgs::LaserScan>("scan/test",5);
        match_scan_pub_ = nh.advertise<visualization_msgs::Marker>(scan_pub_topic,5);
        localization_result_pub_ = nh.advertise<geometry_msgs::PoseStamped>(localization_result_topic,5);
        custom_map_ptr_.reset(new CustomMap(nh,map_topic));
        init_=  false;
        successful_match_ = true;
        init_reloc_pose();
    }


private:
    void scan_callback(const sensor_msgs::LaserScanConstPtr& msg)
    {
        if(!init_)
            return ;
        laser_msg_ = *msg;
        if(successful_match_)
            pointcloud_alignment(laser_msg_,current_pose);
        else{
            std::cout<<"wait for reloc"<<std::endl;
            pointcloud_alignment(laser_msg_,reloc_pose_,1.5);
        }
    }

    void reloc_callback(const geometry_msgs::PoseWithCovarianceStampedConstPtr& msg)
    {
        Eigen::Matrix4d reloc_mat;
        auto pose = msg->pose.pose;
        transform::Pose2Mat4d(reloc_mat,pose);
        pointcloud_alignment(laser_msg_,reloc_mat,1.0);
        current_pose = reloc_mat;
    }

    void odom_callback(const nav_msgs::OdometryConstPtr& msg)
    {
        if(!init_)
        {
            init_ = true;
            auto pose = msg->pose.pose;
            transform::Pose2Mat4d(current_pose,pose);
        }
    }

    bool find_nearest_obstacle(std::vector<Eigen::Vector4d>& hit_vector,std::vector<geometry_msgs::Point>& laser_hit,std::vector<geometry_msgs::Point>& near_obstacle,double threshold)
    {
        int laser_size = hit_vector.size();
        int outlier_size = 0;
        for(int i=0;i<hit_vector.size();i++)
        {
            double obs_x,obs_y;
            double distance = custom_map_ptr_->get_nearest_obstacle(hit_vector[i](0),hit_vector[i](1),obs_x,obs_y);
            if(distance<threshold)
            {
                geometry_msgs::Point p;
                p.x = hit_vector[i](0);
                p.y = hit_vector[i](1);
                laser_hit.emplace_back(p);
                p.x = obs_x;
                p.y = obs_y;
                near_obstacle.emplace_back(p);
            }
            else{
                outlier_size++;
                // std::cout<<distance<<" "<<hit_vector[i](0)<<" "<<hit_vector[i](1)<<" "<<obs_x<<" "<<obs_y<<std::endl;
            }
        }
        visulize_scan_hit(laser_hit);
        double outlier_percent = outlier_size/double(laser_size);
        std::cout<<"outlier_size:"<<outlier_percent<<std::endl;

        if(outlier_percent>0.8 || laser_hit.size()<20)
        {
            return false;
        }
        return true;
    }

    void visulize_scan_hit(std::vector<geometry_msgs::Point>& hit_vector)
    {
        visualization_msgs::Marker marker;
        marker.header.frame_id = "/map";
        marker.action = visualization_msgs::Marker::ADD;
        marker.type = visualization_msgs::Marker::POINTS;
        marker.scale.x = 0.2;
        marker.scale.y = 0.2;
        std_msgs::ColorRGBA color;
        color.a = 1.0;
        color.g = 1.0;
        for(int i=0;i<hit_vector.size();i++)
        {
            marker.points.emplace_back(hit_vector[i]);
            marker.colors.emplace_back(color);
        }
        match_scan_pub_.publish(marker);
    }

    void optimization(std::vector<geometry_msgs::Point>& laser_hit_vector,std::vector<geometry_msgs::Point>& near_obs_vector,Eigen::Vector3d& iter_pose)
    {
        Eigen::MatrixXd laser_hit = Eigen::MatrixXd::Zero(laser_hit_vector.size(),2);
        Eigen::MatrixXd near_obs = Eigen::MatrixXd::Zero(laser_hit_vector.size(),2);
        for(int i=0;i<laser_hit_vector.size();i++)
        {
            laser_hit(i,0) = laser_hit_vector[i].x;
            near_obs(i,0) = near_obs_vector[i].x;
            laser_hit(i,1) = laser_hit_vector[i].y;
            near_obs(i,1) = near_obs_vector[i].y;
        }
        Eigen::MatrixXd laser_avg = Eigen::MatrixXd::Zero(2,1);
        Eigen::MatrixXd obs_avg = Eigen::MatrixXd::Zero(2,1);

        laser_avg(0) = laser_hit.col(0).sum();
        laser_avg(1) = laser_hit.col(1).sum();
        obs_avg(0) = near_obs.col(0).sum();
        obs_avg(1) = near_obs.col(1).sum();

        // for(int i=0;i<laser_hit.rows();i++)
        // {
        //     std::cout<<"i:"<<i<<"    laser_x:"<<laser_hit(i,0)<<"   laser_y"<<laser_hit(i,1)<<"    Obs_x:"<<near_obs(i,0)<<"    Obs_y:"<<near_obs(i,1)<<std::endl;
        // }

        laser_avg /= laser_hit_vector.size();
        obs_avg /= laser_hit_vector.size();

        Eigen::Matrix2d W;
        W.setZero();
        for(int i=0;i<laser_hit_vector.size();i++)
        {
            W +=(laser_hit.row(i).transpose() - laser_avg)*((near_obs.row(i).transpose() - obs_avg).transpose());
        }

        W /= laser_hit_vector.size();

        Eigen::Matrix2d WWT = W*W.transpose();
        Eigen::JacobiSVD<Eigen::MatrixXd> svd(WWT, Eigen::ComputeThinU | Eigen::ComputeThinV );
        Eigen::Matrix2d V = svd.matrixV(), U = svd.matrixU();
        Eigen::Matrix2d  S = U.inverse() * WWT * V.transpose().inverse(); // S = U^-1 * A * VT * -1
        Eigen::Matrix2d S_sqrt = Eigen::Matrix2d::Zero();
        S_sqrt(0,0) = sqrt(S(0,0));
        S_sqrt(1,1) = sqrt(S(1,1));

        Eigen::Matrix2d WWTsqrt = V*S_sqrt*V.transpose();

        Eigen::Matrix2d C = WWTsqrt.inverse()*W;
        Eigen::MatrixXd r(2,1);

        r = obs_avg - C.transpose()*laser_avg;
        iter_pose(0) += r(0);
        iter_pose(1) += r(1);
        iter_pose(2) -= asin(C(1,0));
    }

    void mulls_optimization(std::vector<geometry_msgs::Point>& laser_hit_vector,std::vector<geometry_msgs::Point>& near_obs_vector,Eigen::Vector3d& iter_pose)
    {
        Eigen::Matrix3d ATPA = Eigen::Matrix3d::Zero();
        Eigen::Vector3d ATPb = Eigen::Vector3d::Zero();

        pt2pt_lls_summation(laser_hit_vector,near_obs_vector,ATPA,ATPb);
        ATPA.coeffRef(3) = ATPA.coeffRef(1);
        ATPA.coeffRef(6) = ATPA.coeffRef(2);
        Eigen::Vector3d unknown_x = ATPA.inverse() * ATPb;
        iter_pose+=unknown_x;
    }

	bool pt2pt_lls_summation(std::vector<geometry_msgs::Point>& laser_hit_vector,std::vector<geometry_msgs::Point>& near_obs_vector,
							 Eigen::Matrix3d &ATPA, Eigen::Vector3d &ATPb,float weight=1.0, float residual_window_size = 0.1)
	{
		for (int i = 0; i <laser_hit_vector.size(); i++)
		{
				float px = laser_hit_vector[i].x;
				float py = laser_hit_vector[i].y;
				float qx = near_obs_vector[i].x;
				float qy = near_obs_vector[i].y;

				float dx = px - qx;
				float dy = py - qy;

				float wx, wy, wz;
				wx = weight;
                wy = wx;
                wz = wx;

				// float dist = std::sqrt(qx * qx + qy * qy + qz * qz);

				// if (dist_weight_or_not)
				// 	wx = wx * get_weight_by_dist_adaptive(dist, iter_num);
				// //wx = wx * get_weight_by_dist(dist);

				// if (residual_weight_or_not)
				// 	wx = wx * get_weight_by_residual(std::sqrt(dx * dx + dy * dy + dz * dz), residual_window_size);
				// if (intensity_weight_or_not)
				// 	wx = wx * get_weight_by_intensity(pi + 0.0001, qi + 0.0001);

				// wy = wx;
				// wz = wx;
                
				// unknown x: [tx ty gama] 

				//    0  1  2  3  4  5
				//    6  7  8  9 10 11
				//   12 13 14 15 16 17
				//   18 19 20 21 22 23
				//   24 25 26 27 28 29
				//   30 31 32 33 34 35

                // 0 1 2
                // 3 4 5
                // 6 7 8

				ATPA.coeffRef(0) += wx;
				ATPA.coeffRef(1) += 0;
				ATPA.coeffRef(2) += (-wx * py);
				ATPA.coeffRef(4) += wy;
				ATPA.coeffRef(5) += wy * px;
				ATPA.coeffRef(8) += wx * py * py + wy * px * px;

				ATPb.coeffRef(0) += (-wx * dx);
				ATPb.coeffRef(1) += (-wy * dy);
				ATPb.coeffRef(2) += wx * py * dx - wy * px * dy;
			}
		return true;
	}

    void pointcloud_alignment(const sensor_msgs::LaserScan& msg,Eigen::Matrix4d& init_pose,double match_threshold=0.5)
    {
        // auto pose_cp = init_pose;
        Eigen::Vector3d iter_pose;
        iter_pose(0) = init_pose(0,3);
        iter_pose(1) = init_pose(1,3);
        // iter_pose(2) = init_pose.block<3,3>(0,0).eulerAngles(2,1,0)(0);
        if(init_pose(0,0)*init_pose(0,0)+init_pose(1,0)*init_pose(1,0)>1e-3)
            iter_pose(2) = atan2(init_pose(1,0),init_pose(0,0));
        else 
            iter_pose(2) = 0;
        for(int i=0;i<iter_times_;i++)
        {
            std::vector<Eigen::Vector4d> scan_hit_vector;
            for(int j=rand()%sample_interval_;j<msg.ranges.size();j=j+sample_interval_)
            {
                double r = msg.ranges[j];
                if(isnan(r) || isinf(r))
                    continue;
                double theta = msg.angle_min + j*msg.angle_increment;
                Eigen::Vector4d scan_hit(r*cos(theta),r*sin(theta),0,1);
                scan_hit = init_pose*scan_hit;
                scan_hit_vector.emplace_back(scan_hit);
            }
            std::vector<geometry_msgs::Point> laser_hit,near_obstacle;
            successful_match_ = find_nearest_obstacle(scan_hit_vector,laser_hit,near_obstacle,match_threshold);
            if(!successful_match_)
            {
                return ;
            }
            mulls_optimization(laser_hit,near_obstacle,iter_pose);
            SE2ToMatrix(iter_pose,init_pose);
       }    
       
        geometry_msgs::PoseStamped localization_result;
        localization_result.header.frame_id = "/map";
        transform::Mat4d2Pose(init_pose,localization_result.pose);
        localization_result_pub_.publish(localization_result);
    }



    void Vector4d2Point(std::vector<Eigen::Vector4d>& mat,std::vector<geometry_msgs::Point>& points)
    {
        for(int i=0;i<mat.size();i++)
        {
            geometry_msgs::Point p;
            p.x = mat[i](0);
            p.y = mat[i](1);
            points.emplace_back(p);
        }
    }

    void SE2ToMatrix(Eigen::Vector3d& iter_pose,Eigen::Matrix4d& mat)
    {
        auto last_mat = mat;
        mat(0,3) = iter_pose(0);
        mat(1,3) = iter_pose(1);
        Eigen::AngleAxisd yawAngle(Eigen::AngleAxisd(iter_pose(2),Eigen::Vector3d::UnitZ()));
        mat.block<3,3>(0,0) = yawAngle.toRotationMatrix();
        if(abs(last_mat(0,3)-mat(0,3)>0.3)||abs(last_mat(1,3)-mat(1,3)>0.3))
        {
            std::cout<<"after:"<<iter_pose(2)<<std::endl;
            std::cout<<mat<<std::endl;
            std::cout<<mat<<std::endl<<std::endl;
        }
    }

    void init_reloc_pose(){
        geometry_msgs::Pose p;
        nh.getParam("relocalization_pose/position/x",p.position.x);
        nh.getParam("relocalization_pose/position/y",p.position.y);
        nh.getParam("relocalization_pose/orientation/x",p.orientation.x);
        nh.getParam("relocalization_pose/orientation/y",p.orientation.y);
        nh.getParam("relocalization_pose/orientation/z",p.orientation.z);
        nh.getParam("relocalization_pose/orientaion/w",p.orientation.w);

        transform::Pose2Mat4d(reloc_pose_,p);
        std::cout<<reloc_pose_<<std::endl;

    }



private:
    ros::NodeHandle nh;
    ros::Subscriber scan_sub_;
    ros::Subscriber odom_sub_;
    ros::Publisher match_scan_pub_;
    ros::Publisher test_scan_pub_;
    ros::Publisher localization_result_pub_;
    ros::Subscriber relocalization_sub_;


    std::shared_ptr<CustomMap> custom_map_ptr_;

    bool init_;
    bool successful_match_;
    Eigen::Matrix4d current_pose;
    Eigen::Matrix4d reloc_pose_;
    int sample_interval_;
    int iter_times_;
    sensor_msgs::LaserScan laser_msg_;
};












int main(int argc,char** argv)
{
    ros::init(argc,argv,"localization_node");
    ros::NodeHandle n("~");
    auto map_match = MapMatch(n);
    ros::spin();
}